# 博物馆文物知识库和知识图谱构建工具

import os
import json
import numpy as np
import jieba
import networkx as nx
import matplotlib.pyplot as plt
import faiss
from paddlenlp.embeddings import TokenEmbedding
from PIL import Image
from io import BytesIO
import matplotlib.font_manager as fm

font_path = './fonts/SimHei.ttf'
fontprop = fm.FontProperties(fname=font_path)
MUSEUM_DOCS_PATH = "./data/museum_docs"
VECTOR_DB_PATH = "./data/vector_db"
KNOWLEDGE_GRAPH_PATH = "./data/museum_kg.json"
VECTOR_INDEX_FILE = os.path.join(VECTOR_DB_PATH, "faiss.index")
DOCS_INFO_FILE = os.path.join(VECTOR_DB_PATH, "documents.json")

os.makedirs(MUSEUM_DOCS_PATH, exist_ok=True)
os.makedirs(VECTOR_DB_PATH, exist_ok=True)
os.makedirs(os.path.dirname(KNOWLEDGE_GRAPH_PATH), exist_ok=True)

# -----------------------------------------------------------------------------
# 1. 文物知识库构建
# -----------------------------------------------------------------------------

# 丰富的文物数据集（文本描述）
MUSEUM_ARTIFACTS = {
    "青铜器": {
        "filename": "青铜器.txt",
        "content": """
青铜器是中国古代文明的重要象征，主要兴盛于商周时期（约公元前1600年-公元前256年）。
青铜器种类丰富，主要有礼器、乐器、武器、工具和生活用具等。
礼器类青铜器包括鼎、簋、尊、爵、鬲、彝、钟等，主要用于祭祀活动和宴饮礼仪。
商代青铜器装饰多以兽面纹（饕餮纹）为主，造型庄重威严，给人以神秘感。
西周青铜器纹饰更加精细，出现了更多几何纹样和写实性的动物纹饰。
春秋战国时期的青铜器造型更为多样，工艺更加精湛，出现了镶嵌、错金等复杂工艺。
中国最著名的青铜器包括司母戊鼎（商代）、毛公鼎（西周）、曾侯乙编钟（战国）等。
青铜器铭文是研究古代历史的重要资料，记录了重要的历史事件、政治制度和人物关系。
青铜冶炼技术在中国早在公元前5000年就已出现，到商代已经相当成熟。
青铜器的制作一般采用范铸法，即先做陶范，然后浇注青铜液，冷却后打磨、装饰而成。
"""
    },
    "兵马俑": {
        "filename": "兵马俑.txt",
        "content": """
兵马俑是秦始皇陵墓的陪葬坑，位于陕西省西安市临潼区，距今已有2200多年历史。
1974年3月，陕西临潼县杨家村农民在打井时偶然发现了兵马俑，震惊世界。
兵马俑分为三个坑，总面积约25000平方米，共出土陶俑8000余件，是世界文化遗产。
一号坑是最大的坑，呈长方形，东西长230米，南北宽62米，坑内主要排列着秦军的主力军——步兵、战车和骑兵。
二号坑呈"凹"字形，面积约6000平方米，出土的兵马俑以骑兵、战车、弩兵和步兵为主，被认为是秦军的军阵部署。
三号坑最小，呈"凸"字形，面积约520平方米，是军队的指挥部，出土的兵马俑多为高级军官。
每个兵马俑高1.8米左右，与真人比例相当，重约200公斤，全部由黄土烧制而成。
每个兵马俑的面部特征、服饰、发型各不相同，有不同的表情和姿态，被称为"世界第八大奇迹"。
兵马俑原本是彩色的，全身涂有红、绿、蓝、白、黑、棕等鲜艳颜料，但出土后彩绘大多脱落。
兵马俑展示了秦代高超的制陶工艺、军事组织和艺术成就，也反映了秦始皇的宏大帝国理想和对来世的期望。
"""
    },
    "敦煌壁画": {
        "filename": "敦煌壁画.txt",
        "content": """
敦煌壁画位于甘肃省敦煌市的莫高窟，是世界上现存规模最大、内容最丰富的佛教艺术宝库。
莫高窟是古丝绸之路上的重要文化遗址，始建于十六国的前秦时期（公元366年），历经十六国、北朝、隋、唐、五代、西夏、元等历代的营建。
莫高窟现存洞窟492个，壁画45000多平方米，彩塑2000多座，是世界文化遗产。
敦煌壁画内容以佛教故事为主，包括佛本生故事、佛传故事、经变故事等，也有大量的供养人、飞天、乐伎和世俗生活场景。
唐代是敦煌艺术的鼎盛时期，壁画风格写实优美，人物形象丰满，色彩艳丽，构图严谨，表现了盛唐时期的开放气象。
敦煌壁画采用了熟壁画法，即在干燥的洞窟壁面上涂抹草泥和细泥，待干后再绘制，壁画颜料主要采用矿物质颜料，如朱砂、石青、石绿等。
1900年，道士王圆箓在莫高窟发现了藏经洞，内有大量珍贵文献和绘画，后被英国、法国等国探险家及日本、俄国等国学者收购流失海外。
敦煌壁画艺术风格随朝代变化而变化，北魏时期线条严谨，隋代造型丰满，唐代最为繁荣，五代造型优美，宋代注重写实，西夏融合藏族风格。
敦煌壁画不仅是佛教艺术的瑰宝，也是研究古代社会生活、服饰、音乐、舞蹈、建筑等的重要资料。
敦煌研究院自成立以来，致力于敦煌文化的保护、研究和弘扬，使用现代科技手段对壁画进行保护和修复。
"""
    },
    "唐三彩": {
        "filename": "唐三彩.txt",
        "content": """
唐三彩是盛行于唐代（618-907年）的一种低温釉陶器，因主要釉色为黄、绿、白三色而得名，是中国古代陶瓷艺术的珍品。
唐三彩最早发现于20世纪初陕西西安一带的唐代墓葬中，后在河南洛阳等地也有大量发现，洛阳和西安是两大主要产地。
唐三彩的胎质一般为白色或淡黄色，质地坚硬，主要采用铅釉，在700-900℃左右的温度下烧制而成。
唐三彩的基本釉色为黄、绿、白，有时也加入蓝、紫、褐、黑等色，色彩艳丽，富于变化，流淌自然。
唐三彩的品种多样，包括人物俑、动物俑、生活用具、建筑模型等，以人物俑和马、骆驼等动物俑最为典型。
唐三彩人物俑塑造了各种社会角色，包括武士、侍女、乐伎、胡人等，生动反映了唐代社会生活和审美情趣。
唐三彩动物俑以马和骆驼最为普遍，形象生动逼真，反映了唐代丝绸之路商贸的繁荣和中西文化交流。
唐三彩不仅用于陪葬，也用于宫廷和贵族的日常生活，是唐代物质文明和精神文明的重要载体。
唐三彩制作工艺精湛，造型生动，釉色艳丽，被誉为"唐代艺术的写照"，体现了唐代社会的开放、繁荣与自信。
现存唐三彩主要收藏于中国国家博物馆、陕西历史博物馆、河南博物院等处，也有部分流失海外，成为世界各大博物馆的珍品。
"""
    },
    "越王勾践剑": {
        "filename": "越王勾践剑.txt",
        "content": """
越王勾践剑是中国春秋晚期（约公元前490年前后）的青铜宝剑，因剑身铭文"越王勾践自作用剑"而得名，是中国国宝级文物。
1965年12月，越王勾践剑在湖北江陵望山1号楚墓出土，现藏于湖北省博物馆。
越王勾践剑全长55.7厘米，柄长8.4厘米，剑身长47.3厘米，剑身宽4.6厘米，剑首外翻卷成圆箍形，重875克。
剑身由青铜制成，含铜量约为80%，锡含量约为18%，还有少量的铅、铁和其他微量元素，是中国古代冶金技术的杰出代表。
越王勾践剑经过2500多年仍然锋利无比，出土时仍能削断头发，剑刃上的黑色花纹为古代锻造的精密工艺，被称为"剑之王者"。
剑身上有暗蓝色菱形暗格花纹，剑格上饰以黑色和蓝色相间的菱形几何纹，制作精美，工艺精湛。
越王勾践是春秋时期越国君主，曾被吴王夫差打败，卧薪尝胆，最终灭吴复国，这把剑可能是他复国后所铸。
越王勾践剑采用了多种先进的青铜冶铸技术，包括合金配比、锻造、磨砺、热处理等，体现了中国古代卓越的冶金工艺。
越王勾践剑不仅是一件精美的艺术品，也是研究古代冶金技术、青铜工艺和金属学的重要实物资料。
这把剑的发现对研究春秋战国时期的历史、文化和冶金技术有着重要价值，同时也是中国古代文明的重要象征。
"""
    },
    "莫高窟": {
        "filename": "莫高窟.txt",
        "content": """
莫高窟位于甘肃省敦煌市东南25公里处的鸣沙山东麓，因山势高峻而得名，俗称千佛洞，是中国四大石窟之一。
莫高窟始建于十六国前秦时期（公元366年），由乐尊僧人开凿第一窟，历经十六国、北朝、隋、唐、五代、西夏、元等历代的营建，形成了宏大的规模。
莫高窟现存洞窟492个，南北长1680米，壁画45000多平方米，彩塑2000多座，是世界上现存规模最大、内容最丰富的佛教艺术圣地。
1900年道士王圆箓在莫高窟北端发现了藏经洞（第17窟），内藏大量文书和绘画，被称为"世纪发现"，但后被英国、法国等国探险家及日本、俄国等国学者收购流失海外。
莫高窟壁画内容丰富，以佛教故事为主，包括佛本生故事、佛传故事、经变故事等，也有大量供养人像和世俗生活场景。
莫高窟彩塑以佛、菩萨、弟子、天王、护法神等为主，造型生动，神态各异，色彩艳丽，是中国古代雕塑艺术的精华。
唐代是莫高窟艺术的鼎盛时期，开凿了最多的洞窟，壁画风格写实优美，人物丰满，色彩艳丽，体现了盛唐的开放气象。
莫高窟的建筑布局体现了古代人的巧思，窟内有前室、主室、甬道等构成，平面常呈中心塔柱窟、方形窟、长方形窟等形式。
1961年，莫高窟被列为第一批全国重点文物保护单位；1987年，被联合国教科文组织列入《世界遗产名录》。
敦煌研究院自1944年成立以来，致力于莫高窟的保护、研究和弘扬，使用数字化等现代科技手段对洞窟进行保护和展示。
"""
    },
    "甲骨文": {
        "filename": "甲骨文.txt",
        "content": """
甲骨文是中国已发现的最早的成熟文字，距今约3300年，主要刻写在龟甲和兽骨上，是商代晚期（约公元前1300年-公元前1046年）的占卜记录。
1899年，清朝学者王懿荣在河南安阳小屯村发现了刻有文字的龟甲，后经考古发掘，已出土甲骨文约16万片，著录单字约4500个。
甲骨文的内容主要是商王对天气、收成、疾病、生育、征伐、狩猎等事的占卜记录，是研究商代历史、社会、宗教等的重要资料。
甲骨文已被破译的有约1500个字，其结构特点是象形为主，兼有会意、形声等造字方法，是中国汉字的源头。
甲骨文多为刀刻，笔画主要是直线和曲线，字形结构尚未完全定型，但已具备了汉字的基本特征，有些字与现代汉字相近。
甲骨文揭示了商代的社会组织、政治制度、宗教信仰、天文历法、祭祀活动等，证实了商代历史的真实存在。
甲骨文出土地安阳殷墟被列入《世界遗产名录》，是中国三大古都之一，也是中国历史上第一个有文献记载并经考古发掘证实的商代晚期都城遗址。
甲骨文的发现填补了中国文字发展史上的空白，是中华文明五千年历史的重要物证，被誉为"20世纪中国最重要的考古发现之一"。
甲骨文研究形成了甲骨学这一专门学科，涉及考古学、文字学、历史学、古文字学等多个领域。
目前，甲骨文主要收藏于中国国家博物馆、中国历史博物馆、台北故宫博物院等处，也有部分流失海外，收藏于英国、美国、日本等国的博物馆。
"""
    },
    "四羊方尊": {
        "filename": "四羊方尊.txt",
        "content": """
四羊方尊是商代晚期（约公元前13世纪-公元前11世纪）的青铜礼器，1938年在湖南宁乡黄材镇出土，因器身四角饰有四只羊而得名。
四羊方尊高53.8厘米，口长和腹长分别为28.8厘米和46.4厘米，重34.2千克，现藏于中国国家博物馆。
四羊方尊是典型的商代青铜礼器，为方形，分口沿、颈、腹、足四部分，器身四角各有一只造型生动的羊。
这些羊呈俯首、侧立状，前腿弯曲，羊角高高扬起，触于方尊口沿，形成了四个蒸汽孔，设计巧妙，构思精妙。
方尊通体铸有复杂的饕餮纹、鳞纹和雷纹，线条流畅，纹饰清晰，体现了商代高超的青铜铸造技术。
四羊方尊是商代青铜器中的精品，也是中国青铜艺术的代表作，体现了商代贵族的审美情趣和艺术水平。
方尊是古代的一种重要礼器，主要用于盛放酒，是祭祀和宴飨活动中的重要用具，体现了商代宗教祭祀的礼仪制度。
四羊方尊采用了范铸法，即先做陶范，然后将熔化的青铜浇注入范内，冷却后取出器物，再进行修磨、装饰。
四羊方尊既是实用的礼器，又是精美的艺术品，其完美的设计和精湛的工艺反映了商代社会的文明高度。
四羊方尊被列为中国首批禁止出国（境）展览文物，是中国的国宝级文物，也是中华民族的文化瑰宝。
"""
    },
    "秦始皇陵兵马俑博物馆": {
        "filename": "秦始皇陵兵马俑博物馆.txt",
        "content": """
秦始皇陵兵马俑博物馆位于陕西省西安市临潼区，建于1975年，是为保护、展示和研究秦兵马俑而建立的专题博物馆。
博物馆占地面积约566亩，由兵马俑坑遗址、陈列大厅、文物保护中心、游客服务中心等部分组成。
博物馆内保存着秦始皇陵陪葬坑出土的8000余件兵马俑，是世界上最大的古代军阵俑坑，被誉为"世界第八大奇迹"。
一号坑是最大的一个俑坑，面积14260平方米，出土陶俑、战马和木质战车约6000件，排列成11列纵队，象征着秦军的主力军。
二号坑呈"凹"字形，面积约6000平方米，俑坑中设想出土不同兵种俑约1400件，包括骑兵、步兵和战车兵，体现了秦军的军阵部署。
三号坑是最小的，呈"凸"字形，面积约520平方米，出土军吏俑和武士俑约68件，被解释为秦军的指挥部。
除了兵马俑外，博物馆还展出了铜车马、武器装备等文物，全面展示了秦代军事、文化、艺术和科技成就。
1987年，秦始皇陵及其兵马俑被联合国教科文组织列入《世界遗产名录》，是中国首批世界文化遗产之一。
博物馆年接待游客数百万人次，是世界上最受欢迎的博物馆之一，也是中国重要的爱国主义教育基地和文化旅游目的地。
博物馆不断引入现代化技术手段进行文物保护、修复和展示，如彩绘保护、胶结固化、3D扫描等，以期更好地保护和传承这一世界文化遗产。
"""
    },
    "汉长信宫灯": {
        "filename": "汉长信宫灯.txt",
        "content": """
汉长信宫灯是西汉时期（约公元前206年-公元9年）的青铜宫廷用灯，1968年在河北省满城县中山靖王刘胜墓出土，是中国出土的古代灯具中最精美的一件。
汉长信宫灯高48厘米，重约16公斤，以青铜铸造，全身鎏金，设计精巧，工艺精湛，被誉为"天下第一灯"，现藏于河北博物院。
长信宫灯的主体是一位跪坐的宫女，宫女右手扶灯盘，左手作遮挡烟烬状，神态端庄，造型生动，体现了汉代写实的艺术风格。
灯的灯盘呈圆形，直径13厘米，中央有一个可以旋转的凤鸟形烟道，灯芯燃烧产生的烟气可通过凤鸟的空心喙部排出，巧妙地解决了烟熏问题。
灯的底座为方形，四角各有一个动物足，底部中空，内有储水槽，可用来盛水装灰，防止火星落下引起火灾，设计非常人性化。
长信宫灯铭文刻在灯的底座上，上面刻有"长信尚浴"四字，表明此灯原为皇宫中"长信宫"里的宫灯，而"尚浴"是掌管沐浴的官员。
汉长信宫灯不仅是一件精美的艺术品，也是研究汉代宫廷生活、青铜工艺和灯具发展史的重要实物资料。
从功能性角度看，长信宫灯集照明、防烟、防火于一体，体现了汉代工匠的智慧和汉代科技的发达。
长信宫灯于2002年被列入首批中国禁止出国（境）展览文物，是中国国宝级文物之一。
作为汉代工艺美术的杰出代表，长信宫灯被选为中国古代科技成就的象征，其图案被印在中国第三套人民币5元纸币的背面。
"""
    },
    "曾侯乙编钟": {
        "filename": "曾侯乙编钟.txt",
        "content": """
曾侯乙编钟是战国早期（约公元前433年）的青铜乐器，1978年在湖北随州曾侯乙墓出土，是中国出土的最大、最完整、音律最全的古代编钟。
曾侯乙编钟共65件，分三层八组悬挂在一个精美的钟架上，总重量达2567公斤，是中华文明的瑰宝，现藏于湖北省博物馆。
编钟按大小和音调分为19组，每组包括1-4个钟，最大的一件重达203.6千克，最小的一件重2.4千克。
曾侯乙编钟音域跨越五个半八度，能演奏出宫、商、角、徵、羽五声和变宫、变徵七音，可以演奏出复杂的古代和现代乐曲。
每个钟都具有双音区特点，敲击钟体的不同部位可发出两个不同的音调，增加了演奏的可能性。
钟体表面铸有精美的纹饰，包括蟠螭纹、雷纹和云纹等，部分钟体上刻有铭文，记录了乐律、调性等音乐理论知识。
曾侯乙编钟的出土填补了中国音乐史的重要空白，对研究中国古代音乐理论、乐器制造和青铜铸造技术具有重要价值。
1978年8月，这套编钟经过2000多年后重新被演奏，奏出了《东方红》等乐曲，震惊世界，被誉为"沉睡的乐器"。
曾侯乙编钟采用了精确的范铸法制作，钟体壁厚均匀，音律准确，体现了战国时期高超的科学技术水平和青铜工艺。
作为国宝级文物，曾侯乙编钟被列入中国禁止出国（境）展览文物，是中华民族古老文明的重要象征。
"""
    },
    "《清明上河图》": {
        "filename": "清明上河图.txt",
        "content": """
《清明上河图》是中国北宋画家张择端创作的一幅著名长卷画，创作于北宋徽宗时期（约1085-1145年），描绘了北宋都城汴京（今河南开封）的繁荣景象。
这幅画长卷全长528.7厘米，高24.8厘米，绢本设色，现藏于北京故宫博物院，是中国十大传世名画之一。
画作以春日汴河两岸的市民生活和都城风光为主题，描绘了各行各业、不同阶层人物的生活场景，是北宋城市生活的百科全书。
画卷分为三段：第一段是郊外田园风光，描绘了乡村田野和小桥流水；第二段是郊区景象，描绘了汴河上的虹桥和船只往来；第三段是城内街市，展现了热闹的商业区和各类店铺。
画中共绘制了550多个人物，40多辆车，20多艘大小船只，30多个建筑，包括商店、酒楼、茶肆、戏台等，场景极为丰富。
著名的"虹桥"是画中的标志性场景，桥上熙熙攘攘，车马行人络绎不绝，桥下船只穿梭，反映了当时交通的繁忙。
《清明上河图》采用了散点透视法，视角灵活多变，近大远小，气势恢宏，将众多场景巧妙地串联起来，形成了一幅完整的城市生活画卷。
画作不仅是艺术珍品，也是研究北宋社会经济、城市生活、建筑风格、服饰特点的重要史料，被誉为"中国版的清明之眼"。
《清明上河图》有多个版本，除了北京故宫博物院的宋代原作外，还有台北故宫博物院藏的仿宋本、辽宁省博物馆藏的清院本等。
作为中国绘画艺术的瑰宝，《清明上河图》影响深远，不仅在中国家喻户晓，在国际上也享有盛誉，是中华文化的重要象征。
"""
    },
    "殷墟甲骨": {
        "filename": "殷墟甲骨.txt",
        "content": """
殷墟甲骨是指在河南安阳殷墟发掘出土的刻有甲骨文的龟甲和兽骨，是商代晚期（公元前1300年-公元前1046年）的占卜记录。
殷墟甲骨最早于1899年被发现，当时河南安阳小屯村村民在挖掘中药"龙骨"时偶然发现了刻有奇怪文字的龟甲和兽骨。
经过一个多世纪的考古发掘，目前已出土甲骨约15万片，著录甲骨文单字约4500个，被破译的约1500个，是中国已知最早的成熟文字系统。
甲骨文的内容主要是商王对天气、收成、疾病、生育、征伐、狩猎等事的占卜记录，也有少量关于祭祀、军事等内容的记载。
甲骨占卜的过程是：在龟甲或兽骨上钻凹槽，然后施加热源使其产生裂纹，根据裂纹的形状进行占卜，最后在甲骨上刻上记录。
殷墟甲骨文的发现证实了商代历史的真实存在，为研究商代的政治制度、宗教信仰、经济生活等提供了第一手资料。
甲骨文是汉字的源头，其结构特点是象形为主，兼有会意、形声等造字方法，有些字与现代汉字的形体非常接近。
殷墟甲骨记录了许多商代的历史事件，包括商王的世系、征伐活动、祭祀仪式等，对补充和修正古代文献记载有重要价值。
甲骨文的研究形成了专门的学科——甲骨学，涉及考古学、文字学、历史学、古文字学等多个领域。
2017年，殷墟甲骨文被列入《世界记忆名录》，成为中国的世界记忆遗产，也是中华文明五千年历史的重要物证。
"""
    }
}

def create_document_files():
    """创建文物知识库文档文件"""
    print("开始创建文物知识库...")
    
    for artifact_name, info in MUSEUM_ARTIFACTS.items():
        file_path = os.path.join(MUSEUM_DOCS_PATH, info["filename"])
        with open(file_path, 'w', encoding='utf-8') as f:
            f.write(info["content"])
    
    print(f"已创建 {len(MUSEUM_ARTIFACTS)} 个文物知识文档 ✅")
    return list(MUSEUM_ARTIFACTS.keys())

# -----------------------------------------------------------------------------
# 2. 文本嵌入与向量数据库构建
# -----------------------------------------------------------------------------

class SimpleEmbedding:
    
    def __init__(self):
        """初始化词向量模型"""
        print("加载PaddleNLP词向量模型...")
        self.embedding = TokenEmbedding("w2v.baidu_encyclopedia.target.word-word.dim300")
        self.embedding_dim = 300  # 词向量的维度
        print("词向量模型加载完成 ✅")
        
    def embed_documents(self, texts):
        """将文档列表转换为嵌入向量"""
        embeddings = []
        
        for text in texts:
            words = list(jieba.cut(text))
            word_embeddings = self.embedding.search(words)
            
            if len(word_embeddings) > 0:
                doc_embedding = np.mean(word_embeddings, axis=0)
            else:
                doc_embedding = np.zeros(self.embedding_dim)
            
            embeddings.append(doc_embedding)
        
        return embeddings
    
    def embed_query(self, text):
        return self.embed_documents([text])[0]

def build_vector_database():
    print("开始构建向量数据库...")
    
    # 加载文档
    documents = []
    for filename in os.listdir(MUSEUM_DOCS_PATH):
        if filename.endswith(".txt"):
            file_path = os.path.join(MUSEUM_DOCS_PATH, filename)
            with open(file_path, 'r', encoding='utf-8') as f:
                content = f.read()
            
            documents.append({
                "id": len(documents),
                "content": content,
                "metadata": {"source": file_path}
            })
    
    if not documents:
        print("错误：没有找到文档，请先创建文档文件")
        return False
    
    embedding_model = SimpleEmbedding()
    
    texts = [doc["content"] for doc in documents]
    
    print("生成文档向量嵌入...")
    embeddings = embedding_model.embed_documents(texts)
    
    print("创建FAISS索引...")
    index = faiss.IndexFlatL2(embedding_model.embedding_dim)
    
    embeddings_np = np.array(embeddings).astype('float32')
    index.add(embeddings_np)
    
    print("保存向量索引...")
    faiss.write_index(index, VECTOR_INDEX_FILE)
    
    with open(DOCS_INFO_FILE, 'w', encoding='utf-8') as f:
        json.dump(documents, f, ensure_ascii=False, indent=2)
    
    print(f"向量数据库构建完成，包含 {len(documents)} 个文档 ✅")
    return True

# -----------------------------------------------------------------------------
# 3. 知识图谱构建
# -----------------------------------------------------------------------------

MUSEUM_KNOWLEDGE_GRAPH = {
    'nodes': [
        {'id': 'artifact1', 'name': '青铜器', 'type': 'artifact', 'properties': {'era': '商周时期', 'material': '青铜'}},
        {'id': 'artifact2', 'name': '兵马俑', 'type': 'artifact', 'properties': {'era': '秦朝', 'material': '陶土'}},
        {'id': 'artifact3', 'name': '敦煌壁画', 'type': 'artifact', 'properties': {'era': '多个朝代', 'material': '壁画'}},
        {'id': 'artifact4', 'name': '唐三彩', 'type': 'artifact', 'properties': {'era': '唐朝', 'material': '陶器'}},
        {'id': 'artifact5', 'name': '越王勾践剑', 'type': 'artifact', 'properties': {'era': '春秋晚期', 'material': '青铜'}},
        {'id': 'artifact6', 'name': '甲骨文', 'type': 'artifact', 'properties': {'era': '商代晚期', 'material': '龟甲兽骨'}},
        {'id': 'artifact7', 'name': '四羊方尊', 'type': 'artifact', 'properties': {'era': '商代晚期', 'material': '青铜'}},
        {'id': 'artifact8', 'name': '汉长信宫灯', 'type': 'artifact', 'properties': {'era': '西汉', 'material': '青铜'}},
        {'id': 'artifact9', 'name': '曾侯乙编钟', 'type': 'artifact', 'properties': {'era': '战国早期', 'material': '青铜'}},
        {'id': 'artifact10', 'name': '清明上河图', 'type': 'artifact', 'properties': {'era': '北宋', 'material': '绢本设色'}},
        
        {'id': 'person1', 'name': '秦始皇', 'type': 'person', 'properties': {'birth': '公元前259年', 'death': '公元前210年', 'role': '秦国君主、秦朝皇帝'}},
        {'id': 'person2', 'name': '张择端', 'type': 'person', 'properties': {'era': '北宋', 'role': '画家'}},
        {'id': 'person3', 'name': '勾践', 'type': 'person', 'properties': {'era': '春秋时期', 'role': '越国君主'}},
        {'id': 'person4', 'name': '王羲之', 'type': 'person', 'properties': {'birth': '307年', 'death': '365年', 'role': '书法家'}},
        {'id': 'person5', 'name': '曾侯乙', 'type': 'person', 'properties': {'era': '战国早期', 'role': '曾国国君'}},
        {'id': 'person6', 'name': '刘胜', 'type': 'person', 'properties': {'era': '西汉', 'role': '中山靖王'}},
        
        {'id': 'era1', 'name': '商代', 'type': 'era', 'properties': {'start': '约公元前1600年', 'end': '约公元前1046年'}},
        {'id': 'era2', 'name': '周代', 'type': 'era', 'properties': {'start': '约公元前1046年', 'end': '约公元前256年'}},
        {'id': 'era3', 'name': '春秋时期', 'type': 'era', 'properties': {'start': '约公元前770年', 'end': '约公元前476年'}},
        {'id': 'era4', 'name': '战国时期', 'type': 'era', 'properties': {'start': '约公元前475年', 'end': '约公元前221年'}},
        {'id': 'era5', 'name': '秦朝', 'type': 'era', 'properties': {'start': '公元前221年', 'end': '公元前206年'}},
        {'id': 'era6', 'name': '汉朝', 'type': 'era', 'properties': {'start': '公元前202年', 'end': '公元220年'}},
        {'id': 'era7', 'name': '唐朝', 'type': 'era', 'properties': {'start': '618年', 'end': '907年'}},
        {'id': 'era8', 'name': '宋朝', 'type': 'era', 'properties': {'start': '960年', 'end': '1279年'}},
        
        {'id': 'collection1', 'name': '秦始皇陵', 'type': 'collection', 'properties': {'location': '陕西西安临潼区'}},
        {'id': 'collection2', 'name': '莫高窟', 'type': 'collection', 'properties': {'location': '甘肃敦煌'}},
        {'id': 'collection3', 'name': '中国国家博物馆', 'type': 'collection', 'properties': {'location': '北京'}},
        {'id': 'collection4', 'name': '故宫博物院', 'type': 'collection', 'properties': {'location': '北京'}},
        {'id': 'collection5', 'name': '河北博物院', 'type': 'collection', 'properties': {'location': '河北石家庄'}},
        {'id': 'collection6', 'name': '湖北省博物馆', 'type': 'collection', 'properties': {'location': '湖北武汉'}},
        {'id': 'collection7', 'name': '殷墟', 'type': 'collection', 'properties': {'location': '河南安阳'}},
        
        {'id': 'category1', 'name': '青铜器类', 'type': 'category', 'properties': {'description': '由青铜铸造的器物'}},
        {'id': 'category2', 'name': '陶器类', 'type': 'category', 'properties': {'description': '由黏土烧制的器物'}},
        {'id': 'category3', 'name': '书画类', 'type': 'category', 'properties': {'description': '绘画和书法作品'}},
        {'id': 'category4', 'name': '玉器类', 'type': 'category', 'properties': {'description': '由玉石制成的器物'}},
        {'id': 'category5', 'name': '石刻类', 'type': 'category', 'properties': {'description': '刻在石头上的作品'}},
        {'id': 'category6', 'name': '乐器类', 'type': 'category', 'properties': {'description': '用于演奏音乐的器具'}},
        {'id': 'category7', 'name': '文字类', 'type': 'category', 'properties': {'description': '文字和文献资料'}}
    ],
    'edges': [
        {'source': 'artifact1', 'target': 'era1', 'type': 'created_in', 'properties': {'detail': '商周时期是青铜器的鼎盛时期'}},
        {'source': 'artifact1', 'target': 'era2', 'type': 'created_in', 'properties': {'detail': '周代青铜器工艺更加精湛'}},
        {'source': 'artifact2', 'target': 'era5', 'type': 'created_in', 'properties': {'detail': '兵马俑是秦始皇陵的陪葬品'}},
        {'source': 'artifact3', 'target': 'era7', 'type': 'flourished_in', 'properties': {'detail': '唐代是敦煌壁画的鼎盛时期'}},
        {'source': 'artifact4', 'target': 'era7', 'type': 'created_in', 'properties': {'detail': '唐三彩是唐代特有的陶器'}},
        {'source': 'artifact5', 'target': 'era3', 'type': 'created_in', 'properties': {'detail': '越王勾践剑铸造于春秋晚期'}},
        {'source': 'artifact6', 'target': 'era1', 'type': 'created_in', 'properties': {'detail': '甲骨文是商代晚期的占卜记录'}},
        {'source': 'artifact7', 'target': 'era1', 'type': 'created_in', 'properties': {'detail': '四羊方尊是商代晚期的礼器'}},
        {'source': 'artifact8', 'target': 'era6', 'type': 'created_in', 'properties': {'detail': '汉长信宫灯是西汉时期的宫廷用灯'}},
        {'source': 'artifact9', 'target': 'era4', 'type': 'created_in', 'properties': {'detail': '曾侯乙编钟铸造于战国早期'}},
        {'source': 'artifact10', 'target': 'era8', 'type': 'created_in', 'properties': {'detail': '清明上河图创作于北宋时期'}},
        
        {'source': 'artifact2', 'target': 'person1', 'type': 'associated_with', 'properties': {'description': '兵马俑是秦始皇陵的陪葬品'}},
        {'source': 'artifact10', 'target': 'person2', 'type': 'created_by', 'properties': {'description': '清明上河图由张择端创作'}},
        {'source': 'artifact5', 'target': 'person3', 'type': 'owned_by', 'properties': {'description': '越王勾践剑为越王勾践所有'}},
        {'source': 'artifact9', 'target': 'person5', 'type': 'owned_by', 'properties': {'description': '曾侯乙编钟为曾侯乙所有'}},
        {'source': 'artifact8', 'target': 'person6', 'type': 'owned_by', 'properties': {'description': '汉长信宫灯出土于中山靖王刘胜墓'}},
        
        {'source': 'artifact2', 'target': 'collection1', 'type': 'found_in', 'properties': {'year': '1974年'}},
        {'source': 'artifact3', 'target': 'collection2', 'type': 'located_in', 'properties': {'description': '敦煌壁画位于莫高窟'}},
        {'source': 'artifact7', 'target': 'collection3', 'type': 'collected_by', 'properties': {'description': '四羊方尊现藏于中国国家博物馆'}},
        {'source': 'artifact10', 'target': 'collection4', 'type': 'collected_by', 'properties': {'description': '清明上河图现藏于北京故宫博物院'}},
        {'source': 'artifact8', 'target': 'collection5', 'type': 'collected_by', 'properties': {'description': '汉长信宫灯现藏于河北博物院'}},
        {'source': 'artifact9', 'target': 'collection6', 'type': 'collected_by', 'properties': {'description': '曾侯乙编钟现藏于湖北省博物馆'}},
        {'source': 'artifact6', 'target': 'collection7', 'type': 'found_in', 'properties': {'year': '1899年后陆续发现'}},
        
        {'source': 'artifact1', 'target': 'category1', 'type': 'belongs_to', 'properties': {}},
        {'source': 'artifact2', 'target': 'category2', 'type': 'belongs_to', 'properties': {}},
        {'source': 'artifact3', 'target': 'category3', 'type': 'belongs_to', 'properties': {}},
        {'source': 'artifact4', 'target': 'category2', 'type': 'belongs_to', 'properties': {}},
        {'source': 'artifact5', 'target': 'category1', 'type': 'belongs_to', 'properties': {}},
        {'source': 'artifact6', 'target': 'category7', 'type': 'belongs_to', 'properties': {}},
        {'source': 'artifact7', 'target': 'category1', 'type': 'belongs_to', 'properties': {}},
        {'source': 'artifact8', 'target': 'category1', 'type': 'belongs_to', 'properties': {}},
        {'source': 'artifact9', 'target': 'category6', 'type': 'belongs_to', 'properties': {}},
        {'source': 'artifact10', 'target': 'category3', 'type': 'belongs_to', 'properties': {}},
        
        {'source': 'person1', 'target': 'era5', 'type': 'lived_in', 'properties': {'role': '皇帝'}},
        {'source': 'person2', 'target': 'era8', 'type': 'lived_in', 'properties': {'role': '画家'}},
        {'source': 'person3', 'target': 'era3', 'type': 'lived_in', 'properties': {'role': '君主'}},
        {'source': 'person5', 'target': 'era4', 'type': 'lived_in', 'properties': {'role': '君主'}},
        {'source': 'person6', 'target': 'era6', 'type': 'lived_in', 'properties': {'role': '王侯'}},
        
        {'source': 'collection1', 'target': 'artifact2', 'type': 'houses', 'properties': {'description': '秦始皇陵是兵马俑的出土地和展示地'}},
        {'source': 'collection2', 'target': 'artifact3', 'type': 'houses', 'properties': {'description': '莫高窟中保存有大量敦煌壁画'}},
        {'source': 'collection3', 'target': 'artifact7', 'type': 'houses', 'properties': {'description': '中国国家博物馆收藏有四羊方尊'}},
        {'source': 'collection4', 'target': 'artifact10', 'type': 'houses', 'properties': {'description': '故宫博物院收藏有清明上河图'}},
        {'source': 'collection5', 'target': 'artifact8', 'type': 'houses', 'properties': {'description': '河北博物院收藏有汉长信宫灯'}},
        {'source': 'collection6', 'target': 'artifact9', 'type': 'houses', 'properties': {'description': '湖北省博物馆收藏有曾侯乙编钟'}},
        {'source': 'collection7', 'target': 'artifact6', 'type': 'houses', 'properties': {'description': '殷墟是甲骨文的出土地'}}
    ]
}

def build_knowledge_graph():
    print("开始构建知识图谱...")
    
    graph = nx.DiGraph()
    
    for node in MUSEUM_KNOWLEDGE_GRAPH['nodes']:
        graph.add_node(
            node['id'],
            name=node['name'],
            type=node['type'],
            properties=node['properties']
        )
    
    for edge in MUSEUM_KNOWLEDGE_GRAPH['edges']:
        graph.add_edge(
            edge['source'],
            edge['target'],
            type=edge['type'],
            properties=edge.get('properties', {})
        )
    
    kg_data = {
        'nodes': [],
        'edges': []
    }
    
    for node_id, node_data in graph.nodes(data=True):
        kg_data['nodes'].append({
            'id': node_id,
            'name': node_data['name'],
            'type': node_data['type'],
            'properties': node_data['properties']
        })
    
    for source, target, edge_data in graph.edges(data=True):
        kg_data['edges'].append({
            'source': source,
            'target': target,
            'type': edge_data['type'],
            'properties': edge_data['properties']
        })
    
    with open(KNOWLEDGE_GRAPH_PATH, 'w', encoding='utf-8') as f:
        json.dump(kg_data, f, ensure_ascii=False, indent=2)
    
    print(f"知识图谱构建完成，包含 {graph.number_of_nodes()} 个节点和 {graph.number_of_edges()} 条边 ✅")
    
    return graph

# -----------------------------------------------------------------------------
# 4. 知识图谱可视化
# -----------------------------------------------------------------------------

def visualize_knowledge_graph(graph=None, output_path="museum_knowledge_graph.png"):
    if graph is None:
        if not os.path.exists(KNOWLEDGE_GRAPH_PATH):
            print("错误：知识图谱文件不存在，请先构建知识图谱")
            return False
        
        print("从文件加载知识图谱...")
        graph = nx.DiGraph()
        
        with open(KNOWLEDGE_GRAPH_PATH, 'r', encoding='utf-8') as f:
            kg_data = json.load(f)
        
        for node in kg_data['nodes']:
            graph.add_node(
                node['id'],
                name=node['name'],
                type=node['type'],
                properties=node.get('properties', {})
            )
        
        for edge in kg_data['edges']:
            graph.add_edge(
                edge['source'],
                edge['target'],
                type=edge['type'],
                properties=edge.get('properties', {})
            )
    
    print("可视化知识图谱...")
    plt.figure(figsize=(20, 16))
    
    node_colors = {
        'artifact': 'lightblue',
        'person': 'lightgreen',
        'era': 'salmon',
        'collection': 'yellow',
        'category': 'violet',
        'default': 'gray'
    }
    
    colors = []
    for _, node_data in graph.nodes(data=True):
        node_type = node_data.get('type', 'default')
        colors.append(node_colors.get(node_type, node_colors['default']))
    
    node_sizes = []
    for node in graph.nodes():
        connections = graph.in_degree(node) + graph.out_degree(node)
        size = 300 + connections * 50
        node_sizes.append(size)
    
    pos = nx.spring_layout(graph, seed=42, k=0.3) 
    
    nx.draw_networkx_nodes(graph, pos, 
                          node_color=colors, 
                          alpha=0.8, 
                          node_size=node_sizes,
                          edgecolors='black',
                          linewidths=1)
    
    nx.draw_networkx_edges(graph, pos, 
                          edge_color='gray', 
                          arrows=True, 
                          arrowsize=15,
                          width=1.0,
                          alpha=0.6)
    
    labels = {node_id: data['name'] for node_id, data in graph.nodes(data=True)}
    nx.draw_networkx_labels(graph, pos, labels=labels, font_size=10)
    
    edge_labels = {}
    for s, t, data in graph.edges(data=True):
        if data['type'] in ['created_by', 'owned_by', 'houses', 'associated_with']:
            edge_labels[(s, t)] = data['type']
    
    nx.draw_networkx_edge_labels(graph, pos, edge_labels=edge_labels, font_size=8)
    
    plt.axis('off')
    plt.tight_layout()
    
    plt.savefig(output_path, dpi=300, bbox_inches='tight')
    plt.close()
    
    print(f"知识图谱可视化已保存为 {output_path} ✅")
    return True

# -----------------------------------------------------------------------------
# 5. 主函数
# -----------------------------------------------------------------------------

def build_museum_data():
    print("\n===== 博物馆文物知识库和知识图谱构建工具 =====\n")
    
    artifact_names = create_document_files()
    
    vector_db_success = build_vector_database()
    
    graph = build_knowledge_graph()
    
    vis_success = visualize_knowledge_graph(graph)
    
    print("\n===== 构建完成 =====")
    print(f"文物知识库：{len(artifact_names)}个文物，位于 {MUSEUM_DOCS_PATH}")
    print(f"向量数据库：{'成功' if vector_db_success else '失败'}，位于 {VECTOR_DB_PATH}")
    print(f"知识图谱：{graph.number_of_nodes()}个节点，{graph.number_of_edges()}条边，位于 {KNOWLEDGE_GRAPH_PATH}")
    print(f"知识图谱可视化：{'成功' if vis_success else '失败'}，位于 museum_knowledge_graph.png")
    
    return {
        "artifact_names": artifact_names,
        "vector_db_success": vector_db_success,
        "knowledge_graph": graph,
        "visualization_success": vis_success
    }

if __name__ == "__main__":
    build_museum_data()